import binascii
import functools
import os
import gm_sm4_encrypt_decrypt
import gm_work_key_manager

def get_random_sm4_key_or_iv_bytes():
    return os.urandom(16)

# 使用SM4密钥加解密SM2的私钥们
def encrypt_sm2_private_key_with_sm4(sm2_private_key):
    """
    改造后兼容比特的SM2私钥输入
    """
    # 判断是否为十六进制格式
    if isinstance(sm2_private_key, str) and all(c in "0123456789abcdefABCDEF" for c in sm2_private_key):
        sm2_private_key_hex = sm2_private_key
    else:
        try:
            # 假设比特格式为字节类型，转换为十六进制
            sm2_private_key_hex = binascii.hexlify(sm2_private_key).decode()
        except Exception as e:
            print(f"转换比特私钥为十六进制时出现错误: {e}")
            return None

    sm4_key = gm_work_key_manager.get_work_key_hex()
    try:
        sm4_iv = get_random_sm4_key_or_iv_bytes()
        encrypted_sm2_private_key = gm_sm4_encrypt_decrypt.sm2_private_key_encrypt(sm4_key, sm2_private_key_hex, sm4_iv)
        return encrypted_sm2_private_key, sm4_iv
    except binascii.Error as e:
        print(f"处理十六进制私钥转换字节类型时出现错误: {e}")
        return None
    except Exception as e:
        print(f"加密过程出现未知错误，详细信息: {e}")
        return None


# 使用SM4密钥解密SM2私钥
def decrypt_sm2_private_key_with_sm4(encrypted_sm2_private_key, sm4_iv):
    """
    此函数的作用是利用SM4密钥对已经加密的SM2私钥进行解密操作。
    要求输入的加密后的SM2私钥（encrypted_sm2_private_key）必须是字节（byte）类型的数据，加密时使用的初始向量（IV）也为byte类型，
    解密完成后，函数会将结果以十六进制（hex）字符串的形式返回。

    :param encrypted_sm2_private_key: 经过加密处理后的SM2私钥，数据类型应为字节类型，该参数由加密过程产生并传入用于解密。
    :param sm4_iv: 加密时所使用的初始向量（IV），同样为字节（byte）类型，需要与加密过程中使用的IV保持一致，以确保正确解密。
    :return: 返回解密后的SM2私钥，其数据类型为十六进制字符串；若在解密过程中出现异常，则返回None。
    """
    # 获取工作密钥（以十六进制字符串形式），用于SM4解密操作，需保证和加密时使用的是同一密钥（或相关联的密钥对）
    sm4_key = gm_work_key_manager.get_work_key_hex()
    try:
        # 调用具体的解密函数（假设gm_sm4_encrypt_decrypt模块中有相应的实现）对加密的SM2私钥进行解密，并将结果转换为十六进制字符串形式
        decrypted_sm2_private_key_hex = gm_sm4_encrypt_decrypt.sm2_private_key_decrypt(sm4_key, encrypted_sm2_private_key, sm4_iv)
        return decrypted_sm2_private_key_hex
    except Exception as e:
        print(f"解密SM2私钥时出现异常: {e}")
        return None


def private_key_crypto_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # 处理位置参数检查并加密内容
        new_args = []
        for arg in args:
            if isinstance(arg, dict):
                if 'encrypted_signature_private_key' in arg:
                    arg['encrypted_signature_private_key'],signature_private_key_iv= encrypt_sm2_private_key_with_sm4(arg['encrypted_signature_private_key'])
                    arg['signature_private_key_iv'] = signature_private_key_iv
                if 'encrypted_encryption_private_key' in arg:
                    arg['encrypted_encryption_private_key'],encryption_private_key_iv= encrypt_sm2_private_key_with_sm4(arg['encrypted_encryption_private_key'])
                    arg['encryption_private_key_iv'] = encryption_private_key_iv
            new_args.append(arg)
        # 处理关键字参数检查并加密内容
        new_kwargs = {}
        for key, value in kwargs.items():
            if isinstance(value, dict):
                if 'encrypted_signature_private_key' in value:
                    value['encrypted_signature_private_key'],signature_private_key_iv= encrypt_sm2_private_key_with_sm4(value['encrypted_signature_private_key'])
                    value['signature_private_key_iv'] = signature_private_key_iv
                if 'encrypted_encryption_private_key' in value:
                    value['encrypted_encryption_private_key'],encryption_private_key_iv= encrypt_sm2_private_key_with_sm4(value['encrypted_encryption_private_key'])
                    value['encryption_private_key_iv'] = encryption_private_key_iv
            new_kwargs[key] = value

        result = func(*new_args, **new_kwargs)

        # 处理输出检查（与原代码逻辑一致，保持不变）
        try:
            if isinstance(result, list):
                for element in result:
                    if isinstance(element, dict):
                        if 'encrypted_signature_private_key' in element:
                            element['encrypted_signature_private_key'] = decrypt_sm2_private_key_with_sm4(element['encrypted_signature_private_key'],element['signature_private_key_iv'])
                        if 'encrypted_encryption_private_key' in element:
                            element['encrypted_encryption_private_key'] = decrypt_sm2_private_key_with_sm4(element['encrypted_encryption_private_key'],element['encryption_private_key_iv'])
            return result
        except Exception as e:
            print(f"出现错误: {e}")
            return {"error": str(e)}
    return wrapper
